
#include "xdns-decoder.h"
#include "xdns-data.h"
#include "xdns-header-info.h"
#include "xdns-protocol.h"
#include "xdns-question-info.h"
#include "xdns-resource-record-info.h"
#include <assert.h>

using namespace std;

XdnsError
XdnsDecoder::DecodeHeader(const XdnsData& data,
                          size_t startOffset,
                          XdnsHeaderInfo& header)
{
    unsigned char headerBuf[XDNS_HEADER_SIZE] = {0};
    if (data.GetData(headerBuf, startOffset, XDNS_HEADER_SIZE) != XDNS_HEADER_SIZE)
    {
        return XE_BadResponse;
    }

    header.m_id = XDNS_GET_HEADER_ID(headerBuf);
    header.m_qr = (XDNS_GET_HEADER_QR(headerBuf) != 0);
    header.m_opcode = XDNS_GET_HEADER_OPCODE(headerBuf);
    header.m_aa = (XDNS_GET_HEADER_AA(headerBuf) != 0);
    header.m_tc = (XDNS_GET_HEADER_TC(headerBuf) != 0);
    header.m_rd = (XDNS_GET_HEADER_RD(headerBuf) != 0);
    header.m_ra = (XDNS_GET_HEADER_RA(headerBuf) != 0);
    header.m_z = XDNS_GET_HEADER_Z(headerBuf);
    header.m_rcode = XDNS_GET_HEADER_RCODE(headerBuf);
    header.m_qdcount = XDNS_GET_HEADER_QDCOUNT(headerBuf);
    header.m_ancount = XDNS_GET_HEADER_ANCOUNT(headerBuf);
    header.m_nscount = XDNS_GET_HEADER_NSCOUNT(headerBuf);
    header.m_arcount = XDNS_GET_HEADER_ARCOUNT(headerBuf);

    return XE_Ok;
}

XdnsError
XdnsDecoder::DecodeQuestion(const XdnsData& data,
                            size_t dataOffset,
                            size_t startOffset,
                            XdnsQuestionInfo& question,
                            size_t& dataSize)
{
    string qname;
    size_t qnameSize = 0;
    XdnsError error = DecodeLabels(data,
                                   dataOffset,
                                   startOffset,
                                   qname,
                                   qnameSize);
    if (error != XE_Ok)
    {
        return error;
    }
    startOffset += qnameSize;

    unsigned char qtypeBuf[4] = {0};
    if (data.GetData(qtypeBuf, startOffset, 4) != 4)
    {
        return XE_BadResponse;
    }

    question.m_qtype = XDNS_GET_QUESTION_TYPE(qtypeBuf);
    question.m_qclass = XDNS_GET_QUESTION_CLASS(qtypeBuf);
    question.m_qname = qname;

    dataSize = qnameSize + 4;

    return XE_Ok;
}

XdnsError
XdnsDecoder::DecodeResourceRecord(const XdnsData& data,
                                  size_t dataOffset,
                                  size_t startOffset,
                                  XdnsResourceRecordInfo& resourceRecord,
                                  size_t& dataSize)
{
    string name;
    size_t nameSize = 0;
    XdnsError error = DecodeLabels(data,
                                   dataOffset,
                                   startOffset,
                                   name,
                                   nameSize);
    if (error != XE_Ok)
    {
        return error;
    }
    startOffset += nameSize;

    unsigned char typeBuf[10] = {0};
    if (data.GetData(typeBuf, startOffset, 10) != 10)
    {
        return XE_BadResponse;
    }

    unsigned char rdlength =XDNS_GET_RR_RDLENGTH(typeBuf);

    resourceRecord.m_type = XDNS_GET_RR_TYPE(typeBuf);
    resourceRecord.m_class = XDNS_GET_RR_CLASS(typeBuf);
    resourceRecord.m_rdlength = rdlength;
    resourceRecord.m_ttl = XDNS_GET_RR_TTL(typeBuf);
    resourceRecord.m_rdataOffset = startOffset + 10;
    resourceRecord.m_name = name;

    dataSize = nameSize + 10 + static_cast<size_t>(rdlength);

    return XE_Ok;
}

XdnsError
XdnsDecoder::DecodeLabels(const XdnsData& data,
                          size_t dataOffset,
                          size_t startOffset,
                          string& value,
                          size_t& dataSize)
{
    XdnsError error = XE_Ok;
    size_t offset = startOffset;
    size_t indirectionCount = 0;
    size_t indirectionDataSize = 0;
    size_t wholeDataSize = data.GetDataSize() - dataOffset;
    unsigned char* labelBuf = NULL;
    string decodedString;
    do
    {
        unsigned char buf0 = 0;
        if (data.GetData(&buf0, offset, 1) != 1)
        {
            error = XE_BadResponse;
            break;
        }
        offset += 1;

        if (0 == buf0)
        {
            size_t decodedSize = decodedString.size();
            if (decodedSize > 0)
            {
                // Remove the temporary period ('.') for the last label.
                assert('.' == decodedString.at(decodedSize - 1));
                decodedString.erase(decodedSize - 1);
            }

            break;
        }
        else if ((buf0 & XDNS_INDIRECTION_MASK) != XDNS_INDIRECTION_MASK)
        {
            size_t labelSize = static_cast<size_t>(buf0);
            labelBuf = new (nothrow) unsigned char[labelSize];
            if (NULL == labelBuf)
            {
                error = XE_NoMem;
                break;
            }

            if (data.GetData(labelBuf, offset, labelSize) != labelSize)
            {
                error = XE_BadResponse;
                break;
            }
            offset += labelSize;

            for (size_t i = 0; i < labelSize; ++i)
            {
                char tempChar = static_cast<char>(labelBuf[i]);
                if (('.' == tempChar) || ('\\' == tempChar))
                {
                    decodedString.push_back('\\');
                }
                decodedString.push_back(tempChar);
            }

            /* Append a period as the label separator, it should be
               removed if this label is the last one.
            */
            decodedString.push_back('.');

            delete [] labelBuf;
            labelBuf = NULL;
        }
        else
        {
            if ((++indirectionCount) > wholeDataSize)
            {
                /* It should be a loop caused by hacking if the count of
                   indirection is more than the the whole data size.
                */
                error = XE_BadResponse;
                break;
            }

            unsigned char buf1 = 0;
            if (data.GetData(&buf1, offset, 1) != 1)
            {
                error = XE_BadResponse;
                break;
            }

            if (1 == indirectionCount)
            {
                indirectionDataSize = offset + 1 - startOffset;
            }

            offset = dataOffset + (((buf0 & (~XDNS_INDIRECTION_MASK)) << 8) | buf1);
        }
    }
    while (1);

    if (labelBuf != NULL)
    {
        delete [] labelBuf;
        labelBuf = NULL;
    }

    if (error != XE_Ok)
    {
        return error;
    }

    value = decodedString;
    if (0 == indirectionCount)
    {
        dataSize = offset - startOffset;
    }
    else
    {
        dataSize = indirectionDataSize;
    }

    return XE_Ok;
}

